#!/usr/bin/env python3

import os
import mmap
import struct
import time
import numpy as np
import cv2
import rospy
import threading
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

class SharedMemoryManager:
    """修复后的共享内存管理器"""
    
    def __init__(self, shm_name="/image_shm", max_size=50 * 1024 * 1024):  # 增加到50MB
        self.shm_name = shm_name
        self.max_size = max_size
        self.lock = threading.Lock()
        self._create_shared_memory()
        self.bridge = CvBridge()
    
    def _create_shared_memory(self):
        """创建共享内存区域"""
        try:
            os.unlink(self.shm_name)
        except FileNotFoundError:
            pass
        
        # 创建共享内存文件
        self.fd = os.open(self.shm_name, os.O_CREAT | os.O_RDWR, 0o666)
        os.ftruncate(self.fd, self.max_size)
        
        # 映射共享内存为可读写
        self.shm = mmap.mmap(self.fd, self.max_size, mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE)
    
    def write_images(self, color_msg, depth_msg):
        """将图像写入共享内存 - 修复版本"""
        with self.lock:
            try:
                # 转换ROS消息为OpenCV图像
                color_img = self.bridge.imgmsg_to_cv2(color_msg, "bgr8")
                depth_img = self.bridge.imgmsg_to_cv2(depth_msg, "passthrough")
                
                # 序列化图像数据
                _, color_data = cv2.imencode('.png', color_img)
                _, depth_data = cv2.imencode('.png', depth_img)
                
                color_bytes = color_data.tobytes()
                depth_bytes = depth_data.tobytes()
                
                # 计算数据大小
                color_size = len(color_bytes)
                depth_size = len(depth_bytes)
                total_size = 21 + color_size + depth_size  # 21字节头部(5字节标志+8字节时间戳+8字节大小信息)
                
                if total_size > self.max_size:
                    rospy.logerr(f"图像数据太大({total_size} > {self.max_size})")
                    return False
                
                # 重新设计内存布局，避免覆盖时间戳
                timestamp = int(time.time() * 1000)  # 毫秒时间戳
                
                self.shm.seek(0)
                # 写入标志 (5字节)
                self.shm.write(b'VALID')
                
                # 写入时间戳 (8字节)
                self.shm.write(struct.pack('Q', timestamp))
                
                # 写入图像大小信息 (8字节)
                self.shm.write(struct.pack('II', color_size, depth_size))
                
                # 写入图像数据
                self.shm.write(color_bytes)
                self.shm.write(depth_bytes)
                
                return True
            except Exception as e:
                rospy.logerr(f"写入共享内存错误: {e}")
                return False
    
    def close(self):
        """关闭共享内存"""
        try:
            self.shm.close()
            os.close(self.fd)
            os.unlink(self.shm_name)
        except:
            pass

class ImageAcquisition:
    def __init__(self):
        self.bridge = CvBridge()
        self.shm_manager = SharedMemoryManager()
        self.color_img = None
        self.depth_img = None
        self.lock = threading.Lock()
        
        # ROS订阅
        rospy.Subscriber("/camera/color/image_raw", Image, self.color_callback)
        rospy.Subscriber("/camera/depth/image_raw", Image, self.depth_callback)
        rospy.loginfo("ImageAcquisition initialized, waiting for images...")
        
    def color_callback(self, msg):
        try:
            with self.lock:
                self.color_img = msg
                self.try_process_images()
        except Exception as e:
            rospy.logerr(f"Color error: {e}")

    def depth_callback(self, msg):
        try:
            with self.lock:
                self.depth_img = msg
                self.try_process_images()
        except Exception as e:
            rospy.logerr(f"Depth error: {e}")

    def try_process_images(self):
        if self.color_img is not None and self.depth_img is not None:
            # 写入共享内存
            success = self.shm_manager.write_images(self.color_img, self.depth_img)
            if success:
                rospy.loginfo_throttle(5, "Image pair written to shared memory")
            
            # 重置
            self.color_img = None
            self.depth_img = None

    def shutdown(self):
        """关闭时清理"""
        self.shm_manager.close()
        rospy.loginfo("ImageAcquisition shutdown")

if __name__ == "__main__":
    rospy.init_node("image_acquisition")
    acquisition = ImageAcquisition()
    rospy.on_shutdown(acquisition.shutdown)
    rospy.spin()